package insurance.udf

import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, StructField, StructType}
import org.apache.spark.sql.{Row, SparkSession}

object udaf2 {

  object CalcLxdDxdDxci extends UserDefinedAggregateFunction {
    //每个轮次有哪几个字段参数进来处理
    override def inputSchema: StructType = new StructType()
      .add(StructField("qx_d", DoubleType))
      .add(StructField("qx_ci", DoubleType))
      .add(StructField("lx_d", DoubleType))
    //定义一个缓存容器，用来更新存储每轮的最新值。
    override def bufferSchema: StructType = new StructType()
      .add(StructField("buffer_dx_d", DoubleType))
      .add(StructField("buffer_dx_ci", DoubleType))
      .add(StructField("buffer_lx_d", DoubleType))
    //整个UDAF函数返回的数据类型，此处我们一次性返回一个对象中包含3个字段。
    override def dataType: DataType = new StructType()
      .add(StructField("dx_d", DoubleType))
      .add(StructField("dx_ci", DoubleType))
      .add(StructField("lx_d", DoubleType))

    override def deterministic: Boolean = true

    override def initialize(buffer: MutableAggregationBuffer): Unit = {
    }
    // 分组后的数据，会一行一行地依次进入这个update方法进行处理
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      //input.get(2)就是本轮次的lx_d字段。
      if (null != input.get(2)) {
        //只有当policy_year=1时，才会进来。
        buffer(0) = input.getDouble(0)
        buffer(1) = input.getDouble(1)
        buffer(2) = input.getDouble(2)
      } else {
        //只有当policy_year>1时，才会进来。
        //从缓存中拿出上一轮次的dx_d
        val last_dx_d = buffer.getDouble(0)
        //从缓存中拿出上一轮次的dx_ci
        val last_dx_ci = buffer.getDouble(1)
        //从缓存中拿出上一轮次的lx_d
        val last_lx_d = buffer.getDouble(2)

        //本轮次的qx_d值
        val this_qx_d = input.getDouble(0)
        //本轮次的qx_ci值
        val this_qx_ci = input.getDouble(1)
        //本轮次的最新lx_d值，
        val new_lx_d = last_lx_d - last_dx_d - last_dx_ci
        val new_dx_d = this_qx_d * new_lx_d
        val new_dx_ci = this_qx_ci * new_lx_d
        //本轮次的最新lx_d值，更新进缓存容器。给下轮次使用
        buffer(2) = new_lx_d
        //本轮次的最新dx_d值，更新进缓存容器。给下轮次使用
        buffer(0) = new_dx_d
        //本轮次的最新dx_ci值，更新进缓存容器。给下轮次使用
        buffer(1) = new_dx_ci
      }
    }

    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    }

    override def evaluate(buffer: Row): Any = {
      //把整个缓存容器对象都返回出去，他内部包含了3个字段值
      buffer
    }
  }
}

